import os

import networkx as nx
import pydot
import tqdm
import shutil

# def clean_dot_file(export_dot_file_paths):
#     for export_dot_file_path in export_dot_file_paths:
#         export_dot_file_dir = os.path.dirname(export_dot_file_path)
#         if os.path.exists(export_dot_file_dir) and export_dot_file_dir.split("/")[-1] != "cpg_dir":
#             print("remove: ", export_dot_file_dir)
#             shutil.rmtree(export_dot_file_dir)



def get_export_dot_file_path(dot_file_root):
    dot_file_paths = []
    all_files = os.walk(dot_file_root)

    for root, dirs, files in all_files:
        for file in files:
            if file == "export.dot":
                dot_file_path = os.path.join(root, file)
                dot_file_paths.append(os.path.abspath(dot_file_path))

    return list(set(dot_file_paths))

def get_all_dot_file_path(cpg_file_root):
    dot_file_paths = []
    all_files = os.walk(cpg_file_root)

    for root, dirs, files in all_files:
        for file in files:
            if file.endswith(".dot") and file != "export.dot":
                dot_file_path = os.path.join(root, file)
                dot_file_paths.append(os.path.abspath(dot_file_path))

    return list(set(dot_file_paths))


def read_cpg_file(cpg_file_path):
    nodes_dict = {}
    edges_dict = {}
    source_file_paths = []
    graph = pydot.graph_from_dot_file(cpg_file_path)[0]
    graph_nodes = graph.get_nodes()
    graph_edges = graph.get_edges()

    method_name = os.path.basename(cpg_file_path).split(".")[0]     # dot文件名，即方法名

    nodes = []
    edges = []          # (start_node_id, end_node_id, edge_type)
    for graph_edge in graph_edges:
        start_node_id = int(graph_edge.obj_dict["points"][0])
        end_node_id = int(graph_edge.obj_dict["points"][1])
        edge_type = graph_edge.obj_dict['attributes']["label"].split(" ")[0]
        edges.append([start_node_id, end_node_id, edge_type])

    for graph_node in graph_nodes:
        node_id = int(graph_node.obj_dic t["name"])
        node_type = graph_node.obj_dict['attributes']["label"]

        if "FILENAME" in graph_node.obj_dict['attributes'].keys():
            if graph_node.obj_dict['attributes']["FILENAME"][1:-1] not in ["<empty>", "<unknown>"]:
                source_file_paths.append(graph_node.obj_dict['attributes']["FILENAME"][1:-1])

        node_code = "<empty>"
        if "CODE" in graph_node.obj_dict['attributes'].keys():
            node_code = graph_node.obj_dict['attributes']["CODE"]
        elif "FULL_NAME" in graph_node.obj_dict['attributes'].keys():
            node_code = graph_node.obj_dict['attributes'].get("FULL_NAME")
        elif "NAME" in graph_node.obj_dict['attributes'].keys():
            node_code = graph_node.obj_dict['attributes'].get("NAME")

        if "LINE_NUMBER" in graph_node.obj_dict['attributes'].keys():
            line_number = int(graph_node.obj_dict['attributes'].get("LINE_NUMBER"))
        else:
            line_number = -1

        node_signature = node_code
        if "FULL_NAME" in graph_node.obj_dict['attributes'].keys():
            node_signature = graph_node.obj_dict['attributes'].get("FULL_NAME")
        elif "NAME" in graph_node.obj_dict['attributes'].keys():
            node_signature = graph_node.obj_dict['attributes'].get("NAME")

        nodes.append([node_id, node_type, node_code, node_signature, line_number])

    nodes_dict[method_name] = nodes
    edges_dict[method_name] = edges

    source_file_paths = list(set(source_file_paths))

    return nodes_dict, edges_dict, source_file_paths


def build_CPG(cpg_file_path: str) -> nx.DiGraph:  # 读取dot文件，构图networkx
    """
    build program dependence graph from code

    Args:
        cpg_file_path (str): cpg file path

    Returns: List(CPG)
    """

    nodes_dict, edges_dict, source_file_paths = read_cpg_file(cpg_file_path)
    CPGs = []
    for method_name in nodes_dict.keys():
        CPG = nx.DiGraph(name=method_name)  # 为每一个方法构图
        CPG.graph["file_path"] = source_file_paths
        nodes = nodes_dict.get(method_name)
        edges = edges_dict.get(method_name)
        row_indexs = []
        for node in nodes:
            (node_id, node_type, node_var, node_signature, row_index) = node
            row_indexs.append(int(row_index))
            CPG.add_node(node_id, node_type=node_type, node_var=node_var, node_signature=node_signature,
                         row_index=int(row_index))

        for edge in edges:
            (start_node_id, end_node_id, edge_type) = edge
            CPG.add_edge(start_node_id, end_node_id, edge_type=edge_type)

        # # set(row_indexs) 函数代码少于5行，或 nodes 节点数少于10个， 应该被过滤了。
        # if len(set(row_indexs)) < 5 or len(nodes) < 10:
        #     continue

        if method_name.startswith("bad"):  # 方法名以bad开头是存在错误代码的方法
            CPG.graph["label"] = 1
        else:
            CPG.graph["label"] = 0  # 不存在错误代码

        CPGs.append(CPG)

    return CPGs

def dump_CPG(CPGs, save_cpg_nertworkx_file_dir):  # 保存CPG图到磁盘
    if not os.path.exists(save_cpg_nertworkx_file_dir):
        os.makedirs(save_cpg_nertworkx_file_dir)

    for CPG in CPGs:
        method_name = CPG.name
        out_path = os.path.join(save_cpg_nertworkx_file_dir, method_name + ".nx")
        # print("save cpg to {}...".format(out_path))
        nx.write_gpickle(CPG, out_path)

    # print('save done!')

def read_dot_file_pydot(export_dot_file_path):
    try:
        save_dot_file_dir = os.path.dirname(export_dot_file_path)

        graph = pydot.graph_from_dot_file(export_dot_file_path)[0]
        nodes = graph.get_nodes()
        edges = graph.get_edges()

        edge_map = {}
        reversed_edge_map = {}
        for edge in edges:
            edge_type = "[label=" + edge.obj_dict["attributes"]["label"] + " ]"
            start_node_id = int(edge.obj_dict["points"][0])
            end_node_id = int(edge.obj_dict["points"][1])
            edge_map.setdefault(start_node_id, {}).setdefault(end_node_id, []).append(edge_type)
            reversed_edge_map.setdefault(end_node_id, {}).setdefault(start_node_id, []).append(edge_type)

        methods_node_map = {}  # 方法名: 方法节点node
        all_node_map = {}  # node_id: node
        bad_method_node = None
        good_method_node = None
        for node in nodes:
            all_node_map[int(node.obj_dict["name"])] = node
            node_attributes = node.obj_dict['attributes']
            if node_attributes["label"] == "METHOD" and "NAME" in node_attributes.keys():
                methods_node_map[node_attributes["NAME"]] = node
                if node_attributes["NAME"] == '"bad"':
                    bad_method_node = node
                if node_attributes["NAME"] == '"good"':
                    good_method_node = node

        if bad_method_node is not None:  # 得到 bad 方法的子图
            bad_method_node_ids = set()
            method_node_id = int(bad_method_node.obj_dict['name'])  # 节点的下标
            bad_method_node_ids.add(method_node_id)
            method_name = bad_method_node.obj_dict.get('attributes').get('NAME')[1:-1]  # 去除""

            method_edge_lines = []  # 每一个method对应的边信息，从reversed_edge_map和edge_map中直接构建出来

            bqueue = list()
            visited = set()
            bqueue.append(method_node_id)
            visited.add(method_node_id)
            while bqueue:  # 后向遍历
                fro = bqueue.pop(0)
                bad_method_node_ids.add(fro)
                if fro in reversed_edge_map.keys():
                    for pred in reversed_edge_map.get(fro).keys():
                        if pred not in visited:
                            visited.add(pred)
                            bqueue.append(pred)
                            for edge_type in reversed_edge_map.get(fro).get(pred):
                                edge_line = "  {} -> {} {}".format(pred, fro, edge_type)
                                method_edge_lines.append(edge_line)

            fqueue = list()
            visited = set()
            fqueue.append(method_node_id)
            visited.add(method_node_id)
            while fqueue:  # 前向遍历
                fro = fqueue.pop(0)
                bad_method_node_ids.add(fro)
                if fro in edge_map.keys():
                    for succ in edge_map.get(fro).keys():
                        if succ not in visited:
                            visited.add(succ)
                            fqueue.append(succ)
                            for edge_type in edge_map.get(fro).get(succ):
                                edge_line = "  {} -> {} {}".format(fro, succ, edge_type)
                                method_edge_lines.append(edge_line)

            method_edge_lines = list(set(method_edge_lines))

            save_method_node_lines = []
            for node in nodes:
                if int(node.obj_dict['name']) in bad_method_node_ids:
                    save_method_node_lines.append(node.to_string())

            save_dot_file_path = os.path.join(save_dot_file_dir, "{}.dot".format(method_name))      # 写入bad
            print(save_dot_file_path)
            with open(save_dot_file_path, "w") as f:
                print("digraph \"{}\" ".format(method_name) + "{", file=f)
                print("\n".join(save_method_node_lines), file=f)
                print("\n".join(method_edge_lines), file=f)
                print("}", file=f)

        good_method_nodes = []
        if good_method_node is not None:  # 得到 good 方法的子图
            method_node_id = int(good_method_node.obj_dict['name'])  # 节点的下标
            for node_id in edge_map.get(method_node_id).keys():        # 一度出边
                if all_node_map.get(node_id).obj_dict['attributes']["label"] == "CALL":
                    good_method_name = all_node_map.get(node_id).obj_dict['attributes']["NAME"]
                    if good_method_name in methods_node_map.keys():
                        good_method_nodes.append(methods_node_map.get(good_method_name))


        for good_start_node in good_method_nodes:
            good_method_node_ids = set()
            method_node_id = int(good_start_node.obj_dict['name'])  # 节点的下标
            good_method_node_ids.add(method_node_id)
            method_name = good_start_node.obj_dict.get('attributes').get('NAME')[1:-1]  # 去除""

            method_edge_lines = []  # 每一个method对应的边信息，从reversed_edge_map和edge_map中直接构建出来

            bqueue = list()
            visited = set()
            bqueue.append(method_node_id)
            visited.add(method_node_id)
            while bqueue:  # 后向遍历
                fro = bqueue.pop(0)
                good_method_node_ids.add(fro)
                if fro in reversed_edge_map.keys():
                    for pred in reversed_edge_map.get(fro).keys():
                        if pred not in visited:
                            visited.add(pred)
                            bqueue.append(pred)
                            for edge_type in reversed_edge_map.get(fro).get(pred):
                                edge_line = "  {} -> {} {}".format(pred, fro, edge_type)
                                method_edge_lines.append(edge_line)

            fqueue = list()
            visited = set()
            fqueue.append(method_node_id)
            visited.add(method_node_id)
            while fqueue:  # 前向遍历
                fro = fqueue.pop(0)
                good_method_node_ids.add(fro)
                if fro in edge_map.keys():
                    for succ in edge_map.get(fro).keys():
                        if succ not in visited:
                            visited.add(succ)
                            fqueue.append(succ)
                            for edge_type in edge_map.get(fro).get(succ):
                                edge_line = "  {} -> {} {}".format(fro, succ, edge_type)
                                method_edge_lines.append(edge_line)

            method_edge_lines = list(set(method_edge_lines))

            save_method_node_lines = []
            for node in nodes:
                if int(node.obj_dict['name']) in good_method_node_ids:
                    save_method_node_lines.append(node.to_string())

            save_dot_file_path = os.path.join(save_dot_file_dir, "{}.dot".format(method_name))  # 写入good
            print(save_dot_file_path)
            with open(save_dot_file_path, "w") as f:
                print("digraph \"{}\" ".format(method_name) + "{", file=f)
                print("\n".join(save_method_node_lines), file=f)
                print("\n".join(method_edge_lines), file=f)
                print("}", file=f)

    except Exception as e:
        print(e)
        print(export_dot_file_path)


if __name__ == "__main__":
    # dot_file_path = "/Users/zhang/code2graph/java/source-code/CWE-191/250718-v1/src/main/java/testcases/CWE191_Integer_Underflow/s05/cpg_dir/export.dot"
    # read_dot_file_pydot(dot_file_path)          # 读取export.dot，并生成bad.dot，和good.dot

    dot_file_root = "/Users/shay/code/research/code_2_graph/data_process/data_raw/CWE-191"
    export_dot_file_paths = get_export_dot_file_path(dot_file_root)

    # clean_dot_file(export_dot_file_paths)

    for export_dot_file_path in tqdm.tqdm(export_dot_file_paths):
        read_dot_file_pydot(export_dot_file_path)  # 读取export.dot，并生成bad.dot，和good.dot

    dot_file_paths = get_all_dot_file_path(dot_file_root)

    for dot_file_path in tqdm.tqdm(dot_file_paths):
        CPGs = build_CPG(dot_file_path)
        save_cpg_nertworkx_file_dir = os.path.dirname(dot_file_path)
        dump_CPG(CPGs, save_cpg_nertworkx_file_dir)
